import os
import re
import subprocess
import sys

from log import WriteLog


class Upgrade(WriteLog):
    def __init__(self, **kwargs):
        super(Upgrade, self).__init__()
        self.ipaddress = kwargs.get('ipaddress', '')
        self.ipaddress_std1 = kwargs.get('ipaddress_std1', '')
        self.ipaddress_std2 = kwargs.get('ipaddress_std2', '')
        self.passwd = kwargs.get('userpasswd', '')
        self.pkg_root_path = kwargs.get('pkg_root_path', '')
        self.db_root_path = kwargs.get('db_root_path', '')
        self.upgrade_user_prefix = kwargs.get('upgrade_user_prefix', '')
        self.befor_version = kwargs.get('befor_version', '')
        self.after_version = kwargs.get('after_version', '')
        self.ftp_path = kwargs.get('upgrade_ftp_path', '')
        self.DATE = subprocess.getoutput('date "+%m%d"')
        self.username = ''
        self.env_path = ''
        self.dn1_path = ''
        self.dbport = ''
        self.pkg_path = ''
        self.up_pkg_path = ''
        self.up_pkg_name = ''
        self.log_path = ''

    def set_params(self, rollback=False):
        self.up_pkg_name = self.ftp_path.split('/')[-1]
        if rollback:
            self.log_path = os.path.join(os.getcwd(), 'log', 'run_log',
                                         'upgrade',
                                         f'{self.befor_version}_up_'
                                         f'{self.after_version}_rollback.log')
        else:
            self.log_path = os.path.join(os.getcwd(), 'log', 'run_log',
                                         'upgrade',
                                         f'{self.befor_version}_up_'
                                         f'{self.after_version}.log')
        self.username = subprocess.getoutput(
            f'ls /home|grep {self.upgrade_user_prefix}')
        self.env_path = f'/home/{self.username}/gaussdb.bashrc'
        self.dn1_path = f'{self.db_root_path}/{self.username}/cluster/dn1'
        cmd = f"cat {self.dn1_path}/postgresql.conf|" \
            f"grep -w 'port ='|awk '{{print $3}}'"
        self.dbport = subprocess.getoutput(cmd)
        self.pkg_path = f'{self.pkg_root_path}/{self.username}/' \
            f'{self.DATE}/pkg'
        self.up_pkg_path = f'{self.pkg_root_path}/{self.username}/' \
            f'{self.DATE}/upgrade_pkg'

    def check_befor_upgrade(self, cm_flag=False):
        self.write_file(self.log_path, '升级前准备与检查')
        self.write_file(self.log_path, '1.获取升级前版本')
        cmd = f'su - {self.username} <<EOF\n' \
            f'source {self.env_path}\n' \
            f'gs_ssh -c "gsql -V"\nEOF'
        self.write_file(self.log_path, cmd)
        res = subprocess.getoutput(cmd)
        self.write_file(self.log_path, res)
        assert res.count(self.befor_version) == 3

        self.write_file(self.log_path, '2.健康检查：使用gs_checkos工具完成操作系统状态检查')
        cmd = f'source {self.env_path};gs_checkos -i A'
        self.write_file(self.log_path, cmd)
        res = subprocess.getoutput(cmd)
        self.write_file(self.log_path, res)
        assert 'Abnormal numbers:0' in res

        self.write_file(self.log_path, '3.检查数据库节点磁盘使用率')
        cmd = f'su - {self.username} <<EOF\n' \
            f'source {self.env_path}\n' \
            f'gs_ssh -c "df -h {self.pkg_root_path}|grep -v Filesystem|' \
            f'awk \'{{{{print \\\\\\$5}}}}\'"\nEOF'
        self.write_file(self.log_path, cmd)
        res = subprocess.getoutput(cmd)
        self.write_file(self.log_path, res)
        regex_res = re.split('[:%]', res)
        used_list = [regex_res[2].strip(), regex_res[4].strip(),
                     regex_res[6].strip()]
        for i in used_list:
            assert int(i) <= 70

        self.write_file(self.log_path, '4.检查数据库状态')
        cmd = f'su - {self.username} <<EOF\n' \
            f'source {self.env_path}\n' \
            f'gs_om -t status --all\nEOF'
        self.write_file(self.log_path, cmd)
        res = subprocess.getoutput(cmd)
        self.write_file(self.log_path, res)
        if cm_flag:
            expect = r'\nstate.*:.*Normal'
        else:
            expect = r'instance_state.*:.*Normal'
        regex_res = re.findall(expect, res)
        assert len(regex_res) == 3

        self.write_file(self.log_path,
                        '5.检查参数enable_stream_replication值，该参数为off时不允许升级')
        cmd = f'su - {self.username} <<EOF\n' \
            f'source {self.env_path}\n' \
            f'gs_guc check -N all -I all -c "enable_stream_replication"\nEOF'
        self.write_file(self.log_path, cmd)
        res = subprocess.getoutput(cmd)
        self.write_file(self.log_path, res)
        flag = 'enable_stream_replication=on' in res
        flag1 = 'The value of parameter enable_stream_replication ' \
                'is same on all instances'
        if not (flag and flag1):
            cmd = f'su - {self.username} <<EOF\n' \
                f'source {self.env_path} && ' \
                f'gs_guc set -N all -I all ' \
                f'-c "enable_stream_replication=on" && ' \
                f'gs_om -t restart\nEOF'
            self.write_file(self.log_path, cmd)
            res = subprocess.getoutput(cmd)
            self.write_file(self.log_path, res)
            assert 'Failed instances: 0' in res

    def do_upgrade(self, arg, cm_flag=False):
        self.write_file(self.log_path, '升级操作')
        self.write_file(self.log_path, '1.创建新包目录')
        cmd = f'mkdir -p {self.up_pkg_path}'
        self.write_file(self.log_path, cmd)
        res = subprocess.getoutput(cmd)
        self.write_file(self.log_path, res)
        assert len(res) == 0

        self.write_file(self.log_path, '2.将需要更新的新包上传至新包目录并解压')
        wget_cmd = f'wget {self.ftp_path}'

        cmd = f'cd {self.up_pkg_path} && rm -rf * && ' \
            f'{wget_cmd} && ' \
            f'tar -zxf {self.up_pkg_name} && ' \
            f'for i in `ls|grep \'tar.gz\'`; do tar -zxf "$i"; done'
        self.write_file(self.log_path, cmd)
        res = subprocess.getoutput(cmd)
        self.write_file(self.log_path, res)
        res = subprocess.getoutput(f'ls {self.up_pkg_path}|wc -l')
        assert int(res) > 1

        self.write_file(self.log_path,
                        '3.进入安装包解压出的script目录下,在就地升级或灰度升级前执行前置脚本gs_preinstall')
        shell_cmd = f'''cd {os.path.join(self.up_pkg_path, 'script')} &&
            expect <<EOF
            set timeout 600
            spawn ./gs_preinstall -U {self.username} -G {self.username} \
            -X {self.pkg_path}/{self.username}.xml \
            --sep-env-file={self.env_path}
            expect {{
                "*(yes/no)?" {{ send "yes\\r";exp_continue }}
                "*assword:" {{ send "{self.passwd}\\r" }}
            }}
            expect {{
                "*(yes/no)?" {{ send "yes\\r";exp_continue }}
                "*assword:" {{ send "{self.passwd}\\r";exp_continue }}
                "*assword:" {{ send "{self.passwd}\\r" }}
                eof
            }}''' + '\nEOF\n'
        self.write_file(self.log_path, shell_cmd)
        res = subprocess.getoutput(shell_cmd)
        self.write_file(self.log_path, res)
        assert 'Preinstallation succeeded' in res

        self.write_file(self.log_path, '4.切换至omm用户, gs_om -V查看版本')
        cmd = f'su - {self.username} << EOF\n' \
            f'source {self.env_path}\n' \
            f'gs_ssh -c "gs_om -V"\nEOF\n'
        self.write_file(self.log_path, cmd)
        res = subprocess.getoutput(cmd)
        self.write_file(self.log_path, res)
        assert res.count(f'openGauss OM {self.after_version}') == 3

        self.write_file(self.log_path, '5.数据库状态正常时，开始升级')
        cmd = f'su - {self.username} <<EOF\n' \
            f'source {self.env_path}\n' \
            f'gs_om -t status --all\nEOF'
        self.write_file(self.log_path, cmd)
        res = subprocess.getoutput(cmd)
        self.write_file(self.log_path, res)
        if cm_flag:
            expect = r'\nstate.*:.*Normal'
        else:
            expect = r'instance_state.*:.*Normal'
        regex_res = re.findall(expect, res)
        assert len(regex_res) == 3

        cmd = f'su - {self.username} <<EOF\n' \
            f'source {self.env_path} && ' \
            f'gs_upgradectl -t auto-upgrade ' \
            f'-X {self.pkg_path}/{self.username}.xml {arg}\nEOF'
        self.write_file(self.log_path, cmd)
        res = subprocess.getoutput(cmd)
        self.write_file(self.log_path, res)
        assert 'commit-upgrade' in res

    def after_upgrade(self, cm_flag=False):
        self.write_file(self.log_path, '升级后验证')
        self.write_file(self.log_path, '1.版本验证')
        cmd = f'su - {self.username} <<EOF\n' \
            f'source {self.env_path}\n' \
            f'gs_ssh -c "gsql -V"\nEOF'
        self.write_file(self.log_path, cmd)
        res = subprocess.getoutput(cmd)
        self.write_file(self.log_path, res)
        assert res.count(self.after_version) == 3

        self.write_file(self.log_path, '2.健康检查：使用gs_checkos工具完成操作系统状态检查')
        cmd = f'source {self.env_path};gs_checkos -i A'
        self.write_file(self.log_path, cmd)
        res = subprocess.getoutput(cmd)
        self.write_file(self.log_path, res)
        assert 'Abnormal numbers:0' in res

        self.write_file(self.log_path, '3.检查数据库状态')
        cmd = f'su - {self.username} <<EOF\n' \
            f'source {self.env_path}\n' \
            f'gs_om -t status --all\nEOF'
        self.write_file(self.log_path, cmd)
        res = subprocess.getoutput(cmd)
        self.write_file(self.log_path, res)
        if cm_flag:
            expect = r'\nstate.*:.*Normal'
        else:
            expect = r'instance_state.*:.*Normal'
        regex_res = re.findall(expect, res)
        assert len(regex_res) == 3

    def commit_upgrage(self):
        self.write_file(self.log_path, '提交升级')
        cmd = f'su - {self.username} <<EOF\n' \
            f'source {self.env_path}\n' \
            f'gs_upgradectl -t commit-upgrade ' \
            f'-X {self.pkg_path}/{self.username}.xml\nEOF'
        self.write_file(self.log_path, cmd)
        res = subprocess.getoutput(cmd)
        self.write_file(self.log_path, res)
        assert 'Commit upgrade succeeded' in res or \
               'Commit binary upgrade succeeded' in res

        self.write_file(self.log_path, '版本验证')
        cmd = f'su - {self.username} <<EOF\n' \
            f'source {self.env_path}\n' \
            f'gs_ssh -c "gaussdb -V"\nEOF'
        self.write_file(self.log_path, cmd)
        res = subprocess.getoutput(cmd)
        self.write_file(self.log_path, res)
        assert res.count(self.after_version) == 3

        self.write_file(self.log_path, '配置验证,待补充')
        self.write_file(self.log_path, '数据验证,待补充')

    def rollback(self, arg, expect):
        # 升级完成，提交升级之前进行回滚
        self.write_file(self.log_path, '升级版本回滚')
        cmd = f'su - {self.username} <<EOF\n' \
            f'source {self.env_path}\n' \
            f'gs_upgradectl -t auto-rollback ' \
            f'-X {self.pkg_path}/{self.username}.xml {arg}\nEOF'
        self.write_file(self.log_path, cmd)
        res = subprocess.getoutput(cmd)
        self.write_file(self.log_path, res)
        assert expect in res


if __name__ == '__main__':
    if len(sys.argv) < 11:
        raise Exception('请检查入参！')
    print(f'***************{sys.argv[0]}开始执行****************')
    up_params = dict()
    up_params['ipaddress'] = sys.argv[1]
    up_params['ipaddress_std1'] = sys.argv[2]
    up_params['ipaddress_std2'] = sys.argv[3]
    up_params['userpasswd'] = sys.argv[4]
    up_params['pkg_root_path'] = sys.argv[5]
    up_params['db_root_path'] = sys.argv[6]
    up_params['upgrade_user_prefix'] = sys.argv[7]
    up_params['befor_version'] = sys.argv[8]
    up_params['after_version'] = sys.argv[9]
    up_params['upgrade_ftp_path'] = sys.argv[10]

    up = Upgrade(**up_params)
    up.check_befor_upgrade()
    up.do_upgrade('--grey')
    up.after_upgrade()
    up.commit_upgrage()
    up.after_upgrade()
    print(f'***************{sys.argv[0]}执行完成****************')
